import pathlib
import pickle
import typing

_C = typing.TypeVar("_C")
_R = typing.TypeVar("_R")


class Checkpoint(typing.Generic[_C, _R]):
    def __init__(self, calculation_state: _C, random_state: _R):
        self._calculation_state = calculation_state
        self._random_state = random_state

    def calculation_state(self) -> _C:
        return self._calculation_state

    def random_state(self) -> _R:
        return self._random_state

    def save(self, path: pathlib.Path) -> None:
        with open(path, mode="wb") as stream:
            pickle.dump({
                "state": self._calculation_state,
                "random": self._random_state
            }, stream)

    @staticmethod
    def load(path: pathlib.Path) -> 'Checkpoint':
        with open(path, mode="rb") as stream:
            # TODO: 使用其他序列化方案替代 pickle
            checkpoint = pickle.load(stream)
            return Checkpoint(
                checkpoint["state"],
                checkpoint["random"]
            )
